from enum import Enum

from tortoise import fields, models


class EnhancementType(str, Enum):
    WITHOUT_CONTRAST = "without_contrast"
    ONE_ZONE_CONTRAST = "one_zone_contrast"
    TWO_AND_MORE_ZONES_CONTRAST = "two_and_more_zones_contrast"


class MedicalStudy(models.Model):
    modality: fields.ForeignKeyRelation["models.Modality"] = fields.ForeignKeyField(
        "models.Modality", related_name="medical_studies"
    )
    enhancement = fields.CharEnumField(EnhancementType)
    labor_intensity = fields.DecimalField(max_digits=10, decimal_places=2)

    class Meta:
        unique_together = (("modality", "enhancement"),)
